import logging
import re
from typing import List, Tuple

from PyPDF2 import PdfReader
from langchain_community.embeddings import DashScopeEmbeddings
from langchain_community.vectorstores import FAISS
from langchain_text_splitters import RecursiveCharacterTextSplitter

from rag.ewe.ewe_rag import save_vectorstore, load_knowledge_base, retrieve_relevant_documents


def extract_text_with_page_numbers(pdf) -> Tuple[str,List[int]]:
    """
    从PDF中提取文本并记录每行文本对应的页码
    参数:
        pdf: PDF文件对象
    返回:
        text: 提取的文本内容
        page_numbers: 每行文本对应的页码列表
    """
    text = ""
    page_numbers = []
    for page_number,page in enumerate(pdf.pages,start=1):
        #从当前页面提取文本内容
        extracted_text = page.extract_text()
        if extracted_text:
            text += extracted_text
            page_numbers.extend([page_number] * len(extracted_text.split("\n")))
        else:
            logging.warning(f"No text found on page {page_number}.")
    #使用正则表达式去除"百度文库  - 好好学习，天天向上"文本
    #这个文本在文档中出现，但对实际内容没有意义，需要清除
    text = re.sub(r"百度文库\s*-\s*好好学习，天天向上\s*", "", text)
    #返回提取的文本内容和对应的页码列表
    return text, page_numbers



def split_text_into_chunks(text: str, page_numbers: List[int]) -> Tuple[List[str], dict]:
    """
    将文本内容按照页码切分成多个块
    参数:
        text: 文本内容
        page_numbers: 每行文本对应的页码列表
    返回:
        chunks: 文本块列表
        page_map: 页码映射表，用于记录每个块对应的页码范围

    """
    text_splitter=RecursiveCharacterTextSplitter(
        separators=["\n\n", "\n", ".", " ", ""],
        chunk_size=512,
        chunk_overlap=0,
        length_function=len
    )

    #使用文本分割器将输入文本分割成多个较小的文本块
    chunks = text_splitter.split_text(text)
    print(f"文本被分割成 {len(chunks)} 个块。")

    page_info={chunk: page_numbers[i] for i,chunk in enumerate(chunks)}
    return chunks, page_info

def create_embeddings_and_vectorstore(chunks: List[str]) -> FAISS:
    """
    将文本块转换为向量并创建向量存储

    参数:
        chunks: 文本块列表

    返回:
        knowledgeBase: FAISS向量存储对象
    """
    # 创建嵌入模型实例，使用阿里百炼平台的文本嵌入模型
    # model: 指定使用的嵌入模型名称
    embeddings = DashScopeEmbeddings(
        model="text-embedding-v2"
    )
    # 使用分割后的文本块和嵌入模型创建FAISS向量数据库
    # FAISS是一种高效的相似性搜索库，用于快速检索相似的文本向量
    knowledgeBase = FAISS.from_texts(chunks, embeddings)
    print("已从文本块创建知识库...")

    return knowledgeBase


def save_vectorstore_and_page_info(
    vectorstore: FAISS, 
    page_info: dict, 
    vectorstore_path: str = "faiss_index", 
    page_info_path: str = "page_info.json"
) -> None:
    """
    保存FAISS向量存储和页码信息到磁盘
    
    参数:
        vectorstore: FAISS向量存储对象
        page_info: 页码信息字典
        vectorstore_path: 向量存储保存路径
        page_info_path: 页码信息保存路径
    """
    import json
    import os
    
    # 确保保存路径的目录存在
    vectorstore_dir = os.path.dirname(vectorstore_path)
    if vectorstore_dir and not os.path.exists(vectorstore_dir):
        os.makedirs(vectorstore_dir)
        
    page_info_dir = os.path.dirname(page_info_path)
    if page_info_dir and not os.path.exists(page_info_dir):
        os.makedirs(page_info_dir)
    
    # 保存FAISS向量存储
    vectorstore.save_local(vectorstore_path)
    print(f"向量存储已保存到: {vectorstore_path}")
    
    # 保存页码信息
    with open(page_info_path, 'w', encoding='utf-8') as f:
        json.dump(page_info, f, ensure_ascii=False, indent=2)
    print(f"页码信息已保存到: {page_info_path}")


def load_vectorstore_and_page_info(
    vectorstore_path: str = "faiss_index",
    page_info_path: str = "page_info.json"
) -> Tuple[FAISS, dict]:
    """
    从磁盘加载FAISS向量存储和页码信息
    
    参数:
        vectorstore_path: 向量存储路径
        page_info_path: 页码信息路径
    
    返回:
        vectorstore: FAISS向量存储对象
        page_info: 页码信息字典
    """
    import json
    from langchain_community.embeddings import DashScopeEmbeddings
    
    # 加载FAISS向量存储
    embeddings = DashScopeEmbeddings(
        model="text-embedding-v2"
    )
    vectorstore = FAISS.load_local(vectorstore_path, embeddings, allow_dangerous_deserialization=True)
    print(f"向量存储已从 {vectorstore_path} 加载")
    
    # 加载页码信息
    with open(page_info_path, 'r', encoding='utf-8') as f:
        page_info = json.load(f)
    print(f"页码信息已从 {page_info_path} 加载")
    
    return vectorstore, page_info


def search_relevant_documents(
    query: str,
    vectorstore: FAISS,
    page_info: dict,
    k: int = 5
) -> List[dict]:
    """
    根据问题检索相关文档片段
    
    参数:
        query: 查询问题
        vectorstore: FAISS向量存储对象
        page_info: 页码信息字典
        k: 返回最相关的文档片段数量
    
    返回:
        包含相关文档片段和页码信息的列表
    """
    # 使用向量存储相似性搜索
    docs = vectorstore.similarity_search_with_score(query, k=k)
    
    # 整合页码信息
    results = []
    for doc, score in docs:
        # 获取文档内容
        content = doc.page_content
        
        # 获取对应的页码信息
        page_number = page_info.get(content, "未知")
        
        results.append({
            "content": content,
            "page_number": page_number,
            "score": float(score)
        })
    
    # 按相似度得分排序（得分越低表示越相似）
    results.sort(key=lambda x: x["score"])
    
    return results


def format_search_results(results: List[dict]) -> str:
    """
    格式化搜索结果为易于阅读的字符串
    
    参数:
        results: 搜索结果列表
    
    返回:
        格式化的字符串
    """
    if not results:
        return "未找到相关文档片段。"
    
    formatted_results = []
    for i, result in enumerate(results, 1):
        formatted_results.append(
            f"片段 {i} (第 {result['page_number']} 页, 相似度得分: {result['score']:.4f}):\n"
            f"{result['content']}\n"
            f"{'-' * 50}"
        )
    
    return "\n".join(formatted_results)


def generate_answer_with_llm(
    relevant_docs,
    query: str,
    model_name: str = "qwen-plus"
) -> tuple[str, str]:
    """
    使用大模型基于相关文档生成问题答案
    
    参数:
        relevant_docs: 相关文档片段列表（Document对象列表）
        query: 用户查询问题
        model_name: 使用的大模型名称
    
    返回:
        生成的答案和计费信息
    """
    import os
    from dashscope import Generation
    from langchain.docstore.document import Document
    
    # 检查API密钥
    api_key = os.getenv("DASHSCOPE_API_KEY")
    if not api_key:
        raise ValueError("请设置DASHSCOPE_API_KEY环境变量")
    
    # 构建上下文
    # 检查relevant_docs中的元素类型
    if relevant_docs and isinstance(relevant_docs[0], Document):
        # 如果是Document对象，使用page_content属性
        context = "\n\n".join([doc.page_content for doc in relevant_docs])
    elif relevant_docs and isinstance(relevant_docs[0], dict):
        # 如果是字典，使用"content"键
        context = "\n\n".join([doc["content"] for doc in relevant_docs])
    else:
        # 其他情况，直接转换为字符串
        context = "\n\n".join([str(doc) for doc in relevant_docs])
    
    prompt = f"""
    请根据以下文档内容回答用户的问题。如果文档内容中没有相关信息，请说明无法根据提供的文档回答该问题。

    文档内容:
    {context}

    用户问题: {query}

    请提供简洁、准确的回答:
    """
    
    # 调用大模型生成答案
    try:
        response = Generation.call(
            api_key=api_key,
            model=model_name,
            prompt=prompt,
            max_tokens=1000,
            temperature=0.1,
            top_p=0.9
        )
        
        if response.status_code == 200:
            # 获取计费信息
            cost_info = f"输入token: {response.usage.input_tokens}, 输出token: {response.usage.output_tokens}"
            return response.output.text, cost_info
        else:
            raise Exception(f"大模型调用失败: {response.message}")
    except Exception as e:
        raise Exception(f"生成答案时出错: {str(e)}")


def rag_query(
    query: str,
    vectorstore: FAISS,
    page_info: dict,
    k: int = 5,
    model_name: str = "qwen-plus"
) -> dict:
    """
    完整的RAG查询流程：检索相关文档并生成答案
    
    参数:
        query: 用户查询问题
        vectorstore: FAISS向量存储对象
        page_info: 页码信息字典
        k: 检索文档数量
        model_name: 使用的大模型名称
    
    返回:
        包含答案和相关文档的字典
    """
    # 检索相关文档
    relevant_docs = search_relevant_documents(query, vectorstore, page_info, k)
    
    # 如果没有找到相关文档
    if not relevant_docs:
        return {
            "answer": "抱歉，未找到与您的问题相关的文档内容。",
            "relevant_docs": [],
            "query": query
        }
    
    # 生成答案
    try:
        answer = generate_answer_with_llm(query, relevant_docs, model_name)
    except Exception as e:
        answer = f"生成答案时出错: {str(e)}"
    
    return {
        "answer": answer,
        "relevant_docs": relevant_docs,
        "query": query
    }

def main():
    """
    主函数：测试每个功能函数
    """
    # 创建全局的嵌入模型实例，确保创建和加载时使用相同的模型
    embeddings = DashScopeEmbeddings(model="text-embedding-v2")

    print("=== 测试PDF文本提取 ===")
    try:
        # 读取PDF文件，使用PyPDF2库创建PdfReader对象
        pdf_reader = PdfReader('./浦发上海浦东发展银行西安分行个金客户经理考核办法.pdf')
        # 调用extract_text_with_page_numbers函数提取PDF文本及对应页码
        text, page_numbers = extract_text_with_page_numbers(pdf_reader)
        print(f"成功提取文本，文本长度: {len(text)} 字符")
        print(f"页码数量: {len(page_numbers)}")
        print(f"前5个页码: {page_numbers[:5] if len(page_numbers) > 5 else page_numbers}")
    except Exception as e:
        print(f"PDF文本提取失败: {e}")
        return

    print("\n=== 测试文本切割 ===")
    try:
        # 文本切割
        chunks, page_info = split_text_into_chunks(text, page_numbers)
        print(f"文本被分割成 {len(chunks)} 个块")
        print(f"第一个文本块预览: {chunks[0][:100] if chunks else '无内容'}...")
    except Exception as e:
        print(f"文本切割失败: {e}")
        return

    print("\n=== 测试文本转向量 ===")
    try:
        # 文本转向量并创建向量存储
        knowledgeBase = create_embeddings_and_vectorstore(chunks)
        # 将页码信息附加到知识库对象上，便于后续检索时追踪来源页码
        knowledgeBase.page_info = page_info
        print("成功创建向量存储")
    except Exception as e:
        print(f"文本转向量失败: {e}")
        return

    print("\n=== 测试向量存储保存 ===")
    try:
        # 保存向量存储（可选）
        save_vectorstore(knowledgeBase, page_info, "./vector_db")
        print("向量存储保存成功")
    except Exception as e:
        print(f"向量存储保存失败: {e}")

    print("\n=== 测试向量存储加载 ===")
    try:
        # 加载向量存储，使用相同的嵌入模型
        loaded_knowledgeBase = load_knowledge_base("./vector_db", embeddings)
        print("向量存储加载成功")
    except Exception as e:
        print(f"向量存储加载失败: {e}")
        # 如果加载失败，使用之前创建的知识库
        loaded_knowledgeBase = knowledgeBase

    print("\n=== 测试文档检索 ===")
    try:
        # 设置查询问题
        query = "客户经理每年评聘申报时间是怎样的？"
        # 根据问题检索相关文档
        docs = retrieve_relevant_documents(loaded_knowledgeBase, query)
        print(f"检索到 {len(docs)} 个相关文档")
        if docs:
            print(f"第一个文档预览: {docs[0].page_content[:100] if docs else '无内容'}...")
        else:
            print("未检索到相关文档")
    except Exception as e:
        print(f"文档检索失败: {e}")
        import traceback
        traceback.print_exc()
        return

    print("\n=== 测试大模型生成答案 ===")
    try:
        # 使用大模型生成答案
        response_text, cost_info = generate_answer_with_llm(docs, query)
        print(cost_info)
        print(f"生成的答案预览: {response_text[:200] if response_text else '无内容'}...")
    except Exception as e:
        print(f"大模型生成答案失败: {e}")
        import traceback
        traceback.print_exc()
        return

    print("\n=== 所有函数测试完成 ===")


if __name__ == '__main__':
    main()